from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from schema.response.baseResponse import CustomException
from config.setting import ERROR_LOG


# 数据库连接
engine = create_async_engine("sqlite+aiosqlite:///./test.db", echo=True)
# 创建会话工厂
async_session = sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)

# 数据库基类
Base = declarative_base()


async def init_models():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


# 上下文管理器 用来管理 数据库会话的生命周期， 确保正常释放 （可以通过在这里实现 sql 的关闭 和回滚 在实际代码中采用 依赖注入 的方式 ）
@asynccontextmanager
async def get_session():
    async with async_session() as session:
        async with session.begin():
            try:
                yield session
            except Exception as e:
                await session.rollback()
                ERROR_LOG.error(e)
                raise e

